import torch
import torch.nn as nn
import torch.nn.functional as F

# 尝试导入lpips，如果失败则提供模拟实现
try:
    import lpips
    LPIPS_AVAILABLE = True
except ImportError:
    print("Warning: lpips not available, using fallback implementation")
    LPIPS_AVAILABLE = False

class ProfessionalDisplayLoss(nn.Module):
    """专业级显示质量损失函数"""
    
    def __init__(self, weights=None):
        super().__init__()
        
        # 默认权重配置（基于工程经验调整）
        self.weights = weights or {
            'pixel': 1.0,           # 像素级精度
            'perceptual': 0.1,      # 感知质量
            'color': 0.3,           # 颜色准确性
            'uniformity': 0.2,      # 均匀性
            'low_gray': 0.15,       # 低灰性能
            'texture': 0.05         # 纹理保持
        }
        
        # 基础损失函数
        self.l1_loss = nn.L1Loss()
        # 初始化LPIPS损失函数或使用备用实现
        if LPIPS_AVAILABLE:
            self.lpips_loss = lpips.LPIPS(net='alex')
        else:
            # 备用LPIPS实现 - 简单的L1损失替代
            self.lpips_loss = nn.L1Loss()
        
        # 注册缓冲区用于统计
        self.register_buffer('loss_history', torch.zeros(6))
    
    def color_accuracy_loss(self, pred, target):
        """颜色准确性损失 - 基于你的色准控制经验"""
        # 转换为LAB颜色空间计算颜色差异
        pred_lab = self.rgb_to_lab(pred)
        target_lab = self.rgb_to_lab(target)
        
        # 计算ΔE颜色差异
        delta_e = torch.sqrt(
            torch.sum((pred_lab - target_lab) ** 2, dim=1)
        )
        
        return torch.mean(delta_e)
    
    def rgb_to_lab(self, rgb):
        """RGB转LAB颜色空间（简化版本）"""
        # 这里应该实现完整的RGB->XYZ->LAB转换
        # 为简化，使用近似转换
        r, g, b = rgb[:, 0:1, :, :], rgb[:, 1:2, :, :], rgb[:, 2:3, :, :]
        
        # 简化转换（实际应该使用完整色彩空间转换）
        L = 0.2126 * r + 0.7152 * g + 0.0722 * b
        a = 0.5 * (r - g) + 0.5
        b = 0.5 * (g - b) + 0.5
        
        return torch.cat([L, a, b], dim=1)
    
    def uniformity_loss(self, pred):
        """均匀性损失 - 基于你的均匀性校正经验"""
        # 计算局部亮度变化
        brightness = torch.mean(pred, dim=1, keepdim=True)  # 转为灰度
        
        # 使用拉普拉斯算子检测亮度变化
        laplacian_kernel = torch.tensor([
            [0, 1, 0],
            [1, -4, 1],
            [0, 1, 0]
        ], dtype=torch.float32).view(1, 1, 3, 3).to(pred.device)
        
        # 计算亮度梯度
        brightness_grad = F.conv2d(brightness, laplacian_kernel, padding=1)
        
        # 均匀性损失：惩罚大的亮度变化
        uniformity_loss = torch.mean(torch.abs(brightness_grad))
        
        return uniformity_loss
    
    def low_gray_performance_loss(self, pred, target):
        """低灰性能损失 - 基于你的低灰校正经验"""
        # 提取低亮度区域（模拟低灰显示）
        target_gray = torch.mean(target, dim=1, keepdim=True)
        low_gray_mask = (target_gray < 0.3).float()  # 低亮度区域
        
        if torch.sum(low_gray_mask) == 0:
            return torch.tensor(0.0).to(pred.device)
        
        # 计算低灰区域的差异
        pred_low = pred * low_gray_mask
        target_low = target * low_gray_mask
        
        low_gray_loss = self.l1_loss(pred_low, target_low)
        
        return low_gray_loss
    
    def texture_preservation_loss(self, pred, target):
        """纹理保持损失"""
        # 使用梯度信息保持纹理细节
        pred_grad_x = pred[:, :, :, 1:] - pred[:, :, :, :-1]
        pred_grad_y = pred[:, :, 1:, :] - pred[:, :, :-1, :]
        
        target_grad_x = target[:, :, :, 1:] - target[:, :, :, :-1]
        target_grad_y = target[:, :, 1:, :] - target[:, :, :-1, :]
        
        texture_loss = (
            F.l1_loss(pred_grad_x, target_grad_x) +
            F.l1_loss(pred_grad_y, target_grad_y)
        )
        
        return texture_loss
    
    def forward(self, pred, target):
        # 计算各项损失
        pixel_loss = self.l1_loss(pred, target)
        perceptual_loss = self.lpips_loss(pred, target).mean()
        color_loss = self.color_accuracy_loss(pred, target)
        uniformity_loss = self.uniformity_loss(pred)
        low_gray_loss = self.low_gray_performance_loss(pred, target)
        texture_loss = self.texture_preservation_loss(pred, target)
        
        # 记录损失历史
        losses = torch.tensor([
            pixel_loss.item(), perceptual_loss.item(), color_loss.item(),
            uniformity_loss.item(), low_gray_loss.item(), texture_loss.item()
        ]).to(pred.device)
        self.loss_history = 0.9 * self.loss_history + 0.1 * losses
        
        # 加权总损失
        total_loss = (
            self.weights['pixel'] * pixel_loss +
            self.weights['perceptual'] * perceptual_loss +
            self.weights['color'] * color_loss +
            self.weights['uniformity'] * uniformity_loss +
            self.weights['low_gray'] * low_gray_loss +
            self.weights['texture'] * texture_loss
        )
        
        # 返回总损失和详细损失信息
        loss_details = {
            'total': total_loss.item(),
            'pixel': pixel_loss.item(),
            'perceptual': perceptual_loss.item(),
            'color': color_loss.item(),
            'uniformity': uniformity_loss.item(),
            'low_gray': low_gray_loss.item(),
            'texture': texture_loss.item()
        }
        
        return total_loss, loss_details

# 测试损失函数
def test_professional_loss():
    """测试专业损失函数"""
    # 创建一个简化版本的ProfessionalDisplayLoss用于测试
    class TestProfessionalDisplayLoss(ProfessionalDisplayLoss):
        def __init__(self):
            super().__init__()
            # 在测试模式下始终使用简单的替代实现
            self.lpips_loss = nn.L1Loss()
    
    loss_fn = TestProfessionalDisplayLoss()
    
    # 测试数据
    pred = torch.rand(2, 3, 64, 64)
    target = torch.rand(2, 3, 64, 64)
    
    total_loss, details = loss_fn(pred, target)
    
    print("专业损失函数测试:")
    for key, value in details.items():
        print(f"  {key}: {value:.4f}")
    
    return loss_fn

if __name__ == "__main__":
    test_professional_loss()